from os import path
import torch
import torch.nn.functional as F
import numpy as np
import os
import cv2
import tifffile
import pandas as pd
import rasterio
from rasterio.windows import Window, transform

from tqdm import tqdm
# from model import HuBMAP_model as Model
# from model import MyUNet as Model
import segmentation_models_pytorch as smp
from skimage.measure import label, regionprops

import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')

# def get_data(path):
#     image = tifffile.imread(path)
#     if image.ndim == 5:
#         image = image[0,0,:,:,:]
#         image = np.transpose(image, (1,2,0))
#     return image

def get_data(path):
    identity = rasterio.Affine(1, 0, 0, 0, 1, 0)
    dataset = rasterio.open(path, transform = identity)
    if dataset.count != 3:
        image = np.zeros((dataset.shape[0],dataset.shape[1], 3), dtype=np.uint8)
        for c, subd in enumerate(dataset.subdatasets):
            reader = rasterio.open(subd)
            image[:,:, c] = reader.read()
            reader.close()
    else:
        image = dataset.read()
        image = np.moveaxis(image, 0, -1)
    dataset.close()
    return image

def get_dataset(path):
    identity = rasterio.Affine(1, 0, 0, 0, 1, 0)
    dataset = rasterio.open(path, transform = identity, num_threads='all_cpus')
    layers = []
    if dataset.count != 3:
        subdatasets = dataset.subdatasets
        if len(subdatasets) > 0:
            for subdataset in subdatasets:
                layers.append(rasterio.open(subdataset))
    return dataset, layers

def make_grid(shape, window=512, min_overlap=128):
    """
        Return Array of size (N,4), where N - number of tiles,
        2nd axis represente slices: x1,x2,y1,y2 
    """
    x, y = shape
    nx = x // (window - min_overlap) + 1
    x1 = np.linspace(0, x, num=nx, endpoint=False, dtype=np.int64)
    x1[-1] = x - window
    x2 = (x1 + window).clip(0, x)
    ny = y // (window - min_overlap) + 1
    y1 = np.linspace(0, y, num=ny, endpoint=False, dtype=np.int64)
    y1[-1] = y - window
    y2 = (y1 + window).clip(0, y)
    slices = np.zeros((nx,ny, 4), dtype=np.int64)
    
    for i in range(nx):
        for j in range(ny):
            slices[i,j] = x1[i], x2[i], y1[j], y2[j]    
    return slices.reshape(nx*ny,4)

def mask2rle(mask):
    ''' takes a 2d boolean numpy array and turns it into a space-delimited RLE string '''
    mask = mask.T.reshape(-1) # make 1D, column-first
    mask = np.pad(mask, 1, mode="constant") # make sure that the 1d mask starts and ends with a 0
    starts = np.nonzero((~mask[:-1]) & mask[1:])[0] # start points
    ends = np.nonzero(mask[:-1] & (~mask[1:]))[0] # end points
    rle = np.empty(2 * starts.size, dtype=int) # interlacing...
    rle[0::2] = starts + 1# ...starts...
    rle[1::2] = ends - starts # ...and lengths
    rle = ' '.join([ str(elem) for elem in rle ]) # turn into space-separated string
    return rle

root = "./"
model_flag = "seg9"
output_folder = "{}/submission/{}".format(root, model_flag)
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# model_path = "./work_dir/512_210118resnest50/best.pth"
# model_path = "./work_dir/new_resnest04/latest.pth"
model_path = "./work_dir/{}/latest.pth".format(model_flag)

# dataroot = "../data/test/"
dataroot = "E:/睡眠分期数据/hubmap-kidney-segmentation/test/"
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
score_threshold = 0.4
window_size = 1024
overlap = window_size//4
input_size = (256, 256)
mean = np.array([0.485, 0.456, 0.406])*255
std = np.array([0.229, 0.224, 0.225])*255

out_info = {}

ENCODER = 'efficientnet-b3'
ENCODER_WEIGHTS = 'imagenet'
CLASSES = ['shen_xiao_qiu']
ACTIVATION = 'sigmoid' # could be None for logits or 'softmax2d' for multicalss segmentation
# create segmentation model with pretrained encoder
model = smp.Unet(
    encoder_name=ENCODER, 
    encoder_weights=None, 
    classes=len(CLASSES), 
    activation=ACTIVATION,
)
# model.encoder.load_state_dict(torch.load("E:\\mmdetection\\data\\adv-efficientnet-b4-44fb3a87.pth"))
model.to(device)

model.load_state_dict(torch.load(model_path)['model'])
model.eval()
filelist = list(filter(lambda fn:fn.endswith(".tiff"), os.listdir(dataroot)))
np.random.shuffle(filelist)
main_bar = tqdm(filelist)
for fn in main_bar:
    # img = get_data(dataroot + fn)
    dataset, layers = get_dataset(dataroot + fn)
    # if dataset.count == 3:
    #     continue
    # print(dataset.shape, len(layers))
    pred_mask = np.zeros(dataset.shape[:2], dtype=np.float16)
    mesh_grid = make_grid(dataset.shape[:2], window_size, overlap)
    # pred_mask = np.zeros(img.shape[:2], dtype=np.float16)
    # mesh_grid = make_grid(img.shape[:2], window_size, overlap)[30:]
    with torch.no_grad():
        sub_bar = tqdm(mesh_grid, leave=False)
        sub_bar.set_description("First epoch pred")
        for (x1, x2, y1, y2) in sub_bar:
            # patch = img[x1:x2, y1:y2]
            if dataset.count == 3:
                patch = dataset.read([1, 2, 3], window=Window.from_slices((x1, x2), (y1, y2)))
                patch = np.moveaxis(patch, 0, -1)
                patch = cv2.resize(patch, input_size)
            else:
                patch = np.zeros((*input_size, 3), dtype=np.uint8)
                for i in range(3):
                    patch[:, :, i] = cv2.resize(layers[i].read(window=Window.from_slices((x1, x2), (y1, y2)))[0], input_size)
            orgin_img = cv2.resize(patch, (window_size, window_size))
            
            # patch = cv2.resize(img[x1:x2, y1:y2], input_size)
            # temp = patch.mean((0, 1))
            # temp = sum(abs(temp - temp.mean()))
            # if temp < 8:
            #     continue
            patch = (patch-mean)/std
            patch = np.transpose(patch, (2,0,1))
            img_slice = torch.FloatTensor(patch[np.newaxis,:,:,:]).to(device)

            pred1 = model(img_slice)

            pred2 = model(torch.flip(img_slice, [0, 3]))
            pred2 = torch.flip(pred2, [3, 0])

            pred3 = model(torch.flip(img_slice, [1, 2]))
            pred3 = torch.flip(pred3, [2, 1])

            pred = (pred1 + pred2 + pred3)/3
            # pred = pred.cpu().numpy()[0, 0].astype(np.float16)
            pred = F.interpolate(pred, size=(window_size, window_size), mode="nearest").cpu().numpy()[0, 0].astype(np.float16)
            if pred_mask[x1:x2, y1:y2].shape == pred.shape:
                if x1>overlap and y1>overlap:
                    pred_mask[x1:x1+overlap, y1:y2] = (pred_mask[x1:x1+overlap, y1:y2] + pred[:overlap, :])/2
                    pred_mask[x1+overlap:x2, y1:y1+overlap] = (pred_mask[x1+overlap:x2, y1:y1+overlap] + pred[overlap:, :overlap])/2
                    pred_mask[x1+overlap:x2, y1+overlap:y2] = pred[overlap:, overlap:]
                else:
                    pred_mask[x1:x2, y1:y2] = pred
                # pred_mask[x1:x2, y1:y2] = pred
            # if x1 < 2048 or y1 < 2048:
            #     continue
            pred = (pred >= score_threshold).astype(np.uint8)
            pred2 = (pred_mask[x1:x2, y1:y2] >= score_threshold).astype(np.uint8)
            maskpointnum = pred.sum()
            if (maskpointnum/(window_size*window_size)) > 0.01:# and np.random.randint(0, 100) > 60:
                # sub_bar.write("{:^5.2f}".format(temp))
                plt.figure(figsize=(15,15))
                plt.subplot("131")
                plt.axis('off')
                # plt.title(temp)
                plt.imshow(orgin_img)
                # plt.imshow(pred,alpha=0.5, cmap='gray')
                plt.imshow(pred,alpha=0.5, cmap='gray')
                plt.subplot("132")
                plt.axis('off')
                # plt.imshow(pred)
                plt.imshow(orgin_img)
                # plt.imshow(pred2,alpha=0.5, cmap='gray')
                plt.subplot("133")
                plt.axis('off')
                plt.imshow(orgin_img)
                # plt.imshow(pred,alpha=0.5, cmap='gray')
                plt.imshow(pred2,alpha=0.5, cmap='gray')
                plt.show()
        sub_bar.clear()
        sub_bar.close()
    
    np.save("{}/{}_rawpred.npy".format(output_folder, fn), pred_mask)

    pred_mask = (pred_mask >= score_threshold).astype(np.uint8)
    out_info[len(out_info)] = {'id':fn.split(".")[0], 'predicted': mask2rle(pred_mask)}

submission = pd.DataFrame.from_dict(out_info, orient='index')
submission.to_csv('{}/submission.csv'.format(output_folder), index=False)